#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys

import numpy as np
from PIL import Image, ImageQt
from PyQt5.QtGui import QColor
from PyQt5.QtWidgets import QMainWindow, QDesktopWidget, QApplication

from util.functions import softmax
from ui.gui import Ui_MainWindow
from cnn.deep_convnet import DeepConvNet


class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        # 初始化UI
        self.network = None

        self.setupUi(self)
        self.move_to_center()
        self.register_event()
        self.create_network()

        self.number = 0.0
        self.confidence = 0.0

    def create_network(self):
        self.network = DeepConvNet()
        self.network.load_params("model/deep_convnet_params.pkl")

    # 移动窗口至屏幕中央
    def move_to_center(self):
        # 获得窗口
        frame = self.frameGeometry()
        # 获得屏幕中心点
        screen_center = QDesktopWidget().availableGeometry().center()
        # 显示到屏幕中心
        frame.moveCenter(screen_center)
        self.move(frame.topLeft())

    def register_event(self):
        self.btn_clear.clicked.connect(self.clear_draw_area)
        self.btn_predict.clicked.connect(self.predict_number)

    # 清除数据
    def clear_draw_area(self):
        self.draw_area.clear_draw_panel()
        self.lable_result.clear()
        self.lable_confidence_value.clear()

    # 识别数字
    def predict_number(self):
        # 获取QImage格式图像
        number_image = self.draw_area.get_draw_panel_image()

        # 转换成pil image类型图像再处理
        pil_image = ImageQt.fromqimage(number_image)
        pil_image = pil_image.resize((28, 28), Image.ANTIALIAS)

        # 转为灰度图，并修改分辨率为 28*28，后归一化
        image = np.array(pil_image.convert('L')).reshape(1, 1, 28, 28) / 255.0
        image = 1 - image  # 取反，绘制出来的数字的像素为黑色，需要是白色

        # 开始识别
        predict_result = self.network.predict(image)

        # 激活函数softmax得出最终识别结果
        predict_result = softmax(predict_result)

        # 预测的数字
        self.number = np.argmax(predict_result)

        # 置信度
        self.confidence = predict_result[0, self.number]

        self.lable_result.setText(
            "<html><head/><body><p align=\"center\"><span style=\" font-size:88pt;\">%d</span></p></body></html>" % self.number)
        self.lable_confidence_value.setText(
            "<html><head/><body><p><span style=\" font-size:12pt;\">%.8f</span></p></body></html>" % self.confidence)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    Gui = MainWindow()
    Gui.show()

    sys.exit(app.exec_())
